// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#pragma once

#include <algorithm>
#include <utility>
#include <vector>


#include <nebula/util/algorithm.h>
#include <turbo/functional/functional.h>
#include <turbo/log/logging.h>

namespace nebula::internal {

    template<typename T>
    std::vector<T> DeleteVectorElement(const std::vector<T> &values, size_t index) {
                DKCHECK(!values.empty());
                DKCHECK_LT(index, values.size());
        std::vector<T> out;
        out.reserve(values.size() - 1);
        for (size_t i = 0; i < index; ++i) {
            out.push_back(values[i]);
        }
        for (size_t i = index + 1; i < values.size(); ++i) {
            out.push_back(values[i]);
        }
        return out;
    }

    template<typename T>
    std::vector<T> AddVectorElement(const std::vector<T> &values, size_t index,
                                    T new_element) {
                DKCHECK_LE(index, values.size());
        std::vector<T> out;
        out.reserve(values.size() + 1);
        for (size_t i = 0; i < index; ++i) {
            out.push_back(values[i]);
        }
        out.emplace_back(std::move(new_element));
        for (size_t i = index; i < values.size(); ++i) {
            out.push_back(values[i]);
        }
        return out;
    }

    template<typename T>
    std::vector<T> ReplaceVectorElement(const std::vector<T> &values, size_t index,
                                        T new_element) {
                DKCHECK_LE(index, values.size());
        std::vector<T> out;
        out.reserve(values.size());
        for (size_t i = 0; i < index; ++i) {
            out.push_back(values[i]);
        }
        out.emplace_back(std::move(new_element));
        for (size_t i = index + 1; i < values.size(); ++i) {
            out.push_back(values[i]);
        }
        return out;
    }

    template<typename T, typename Predicate>
    std::vector<T> FilterVector(std::vector<T> values, Predicate &&predicate) {
        auto new_end = std::remove_if(values.begin(), values.end(),
                                      [&](const T &value) { return !predicate(value); });
        values.erase(new_end, values.end());
        return values;
    }

    template<typename Fn, typename From,
            typename To = decltype(std::declval<Fn>()(std::declval<From>()))>
    std::vector<To> MapVector(Fn &&map, const std::vector<From> &source) {
        std::vector<To> out;
        out.reserve(source.size());
        std::transform(source.begin(), source.end(), std::back_inserter(out),
                       std::forward<Fn>(map));
        return out;
    }

    template<typename Fn, typename From,
            typename To = decltype(std::declval<Fn>()(std::declval<From>()))>
    std::vector<To> MapVector(Fn &&map, std::vector<From> &&source) {
        std::vector<To> out;
        out.reserve(source.size());
        std::transform(std::make_move_iterator(source.begin()),
                       std::make_move_iterator(source.end()), std::back_inserter(out),
                       std::forward<Fn>(map));
        return out;
    }

    /// \brief Like MapVector, but where the function can fail.
    template<typename Fn, typename From = turbo::call_traits::argument_type<0, Fn>,
            typename To = typename turbo::call_traits::return_type<Fn>::value_type>
    turbo::Result<std::vector<To>> MaybeMapVector(Fn &&map, const std::vector<From> &source) {
        std::vector<To> out;
        out.reserve(source.size());
        TURBO_RETURN_NOT_OK(MaybeTransform(source.begin(), source.end(),
                                            std::back_inserter(out), std::forward<Fn>(map)));
        return out;
    }

    template<typename Fn, typename From = turbo::call_traits::argument_type<0, Fn>,
            typename To = typename turbo::call_traits::return_type<Fn>::value_type>
    turbo::Result<std::vector<To>> MaybeMapVector(Fn &&map, std::vector<From> &&source) {
        std::vector<To> out;
        out.reserve(source.size());
        TURBO_RETURN_NOT_OK(MaybeTransform(std::make_move_iterator(source.begin()),
                                            std::make_move_iterator(source.end()),
                                            std::back_inserter(out), std::forward<Fn>(map)));
        return std::move(out);
    }

    template<typename T>
    std::vector<T> FlattenVectors(const std::vector<std::vector<T>> &vecs) {
        std::size_t sum = 0;
        for (const auto &vec: vecs) {
            sum += vec.size();
        }
        std::vector<T> out;
        out.reserve(sum);
        for (const auto &vec: vecs) {
            out.insert(out.end(), vec.begin(), vec.end());
        }
        return out;
    }

    template<typename T>
    turbo::Result<std::vector<T>> UnwrapOrRaise(std::vector<turbo::Result<T>> &&results) {
        std::vector<T> out;
        out.reserve(results.size());
        auto end = std::make_move_iterator(results.end());
        for (auto it = std::make_move_iterator(results.begin()); it != end; it++) {
            if (!it->ok()) {
                return it->status();
            }
            out.push_back(it->move_value_unsafe());
        }
        return out;
    }

    template<typename T>
    turbo::Result<std::vector<T>> UnwrapOrRaise(const std::vector<turbo::Result<T>> &results) {
        std::vector<T> out;
        out.reserve(results.size());
        for (const auto &result: results) {
            if (!result.ok()) {
                return result.status();
            }
            out.push_back(result.value_unsafe());
        }
        return out;
    }

}  // namespace nebula::internal
